from pickletools import optimize
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from modelzoo import Net
from mymodel import DesNet, ResNet
import os

if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(device)

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    batch_size = 4
    root = '/root/autodl-tmp/dataset/cifar-10-python' if torch.cuda.is_available() else 'E:\\MyBackups\\files\\Datasets\\cifar-10-python'

    trainset = torchvision.datasets.CIFAR10(root=root, train=True,
                                            download=False, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                              shuffle=True, num_workers=2)

    DIR = 'DenseNet'

    if not os.path.exists(DIR):
        os.mkdir(DIR)
    net = DesNet()

    net.to(device)
    # print(net)
    criterion = nn.CrossEntropyLoss()
    # optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    optimizer = optim.Adam(net.parameters(), lr=0.001)
    loss_list = []
    acc_list = []
    for epoch in range(2):  # loop over the dataset multiple times

        running_loss = 0.0
        running_nacc = 0
        num = 0
        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data
            inputs = inputs.to(device=device)
            labels = labels.to(device=device)
            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            running_nacc += torch.sum(torch.argmax(outputs.data, dim=1) == labels).item()
            num += len(labels)
            if i % 2000 == 1999:  # print every 2000 mini-batches
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / num:.3f} acc: {running_nacc / num:.3f}')
                loss_list.append(running_loss / num)
                acc_list.append(running_nacc / num)
                running_loss = 0.0
                running_nacc = 0
                num = 0

    print('Finished Training')

    torch.save(net.state_dict(), os.path.join(DIR, "model.pth"))
    torch.save(dict(train_loss=loss_list, train_acc=acc_list), os.path.join(DIR, "train_hist.pth"))
